# -*- coding: utf-8 -*-

###########################################################################
#
#   Copyright (c) 2020  Bresee, Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
###########################################################################

# pgsql operation

# from collections.abc import Callable
# from typing import Any, TypeVar, overload
# import psycopg2

# from collections import OrderedDict
# from collections.abc import Callable
# from typing import Any, NamedTuple, TypeVar, overload

# from psycopg2._ipaddress import register_ipaddress as register_ipaddress
# from psycopg2._json import (
#     Json as Json,
#     register_default_json as register_default_json,
#     register_default_jsonb as register_default_jsonb,
#     register_json as register_json,
# )
# from psycopg2._psycopg import (
#     REPLICATION_LOGICAL as REPLICATION_LOGICAL,
#     REPLICATION_PHYSICAL as REPLICATION_PHYSICAL,
#     ReplicationConnection as _replicationConnection,
#     ReplicationCursor as _replicationCursor,
#     ReplicationMessage as ReplicationMessage,
# )
# from psycopg2._range import (
#     DateRange as DateRange,
#     DateTimeRange as DateTimeRange,
#     DateTimeTZRange as DateTimeTZRange,
#     NumericRange as NumericRange,
#     Range as Range,
#     RangeAdapter as RangeAdapter,
#     RangeCaster as RangeCaster,
#     register_range as register_range,
# )

# #from .extensions import connection as _connection, cursor as _cursor, quote_ident as quote_ident
# _cursor = {}
# _T_cur = TypeVar("_T_cur", bound=_cursor)


# class DictCursorBase(_cursor):
#     def __init__(self, *args, **kwargs) -> None: ...

# class DictConnection(_connection):
#     @overload
#     def cursor(self, name: str | bytes | None = ..., *, withhold: bool = ..., scrollable: bool | None = ...) -> DictCursor: ...
#     @overload
#     def cursor(
#         self,
#         name: str | bytes | None = ...,
#         *,
#         cursor_factory: Callable[..., _T_cur],
#         withhold: bool = ...,
#         scrollable: bool | None = ...,
#     ) -> _T_cur: ...
#     @overload
#     def cursor(
#         self, name: str | bytes | None, cursor_factory: Callable[..., _T_cur], withhold: bool = ..., scrollable: bool | None = ...
#     ) -> _T_cur: ...

# class DictCursor(DictCursorBase):
#     def __init__(self, *args, **kwargs) -> None: ...
#     index: Any
#     def execute(self, query, vars: Any | None = ...): ...
#     def callproc(self, procname, vars: Any | None = ...): ...
#     def fetchone(self) -> DictRow | None: ...  # type: ignore[override]
#     def fetchmany(self, size: int | None = ...) -> list[DictRow]: ...  # type: ignore[override]
#     def fetchall(self) -> list[DictRow]: ...  # type: ignore[override]
#     def __next__(self) -> DictRow: ...  # type: ignore[override]

# class DictRow(list[Any]):
#     def __init__(self, cursor) -> None: ...
#     def __getitem__(self, x): ...
#     def __setitem__(self, x, v) -> None: ...
#     def items(self): ...
#     def keys(self): ...
#     def values(self): ...
#     def get(self, x, default: Any | None = ...): ...
#     def copy(self): ...
#     def __contains__(self, x): ...
#     def __reduce__(self): ...

# class RealDictConnection(_connection):
#     @overload
#     def cursor(
#         self, name: str | bytes | None = ..., *, withhold: bool = ..., scrollable: bool | None = ...
#     ) -> RealDictCursor: ...
#     @overload
#     def cursor(
#         self,
#         name: str | bytes | None = ...,
#         *,
#         cursor_factory: Callable[..., _T_cur],
#         withhold: bool = ...,
#         scrollable: bool | None = ...,
#     ) -> _T_cur: ...
#     @overload
#     def cursor(
#         self, name: str | bytes | None, cursor_factory: Callable[..., _T_cur], withhold: bool = ..., scrollable: bool | None = ...
#     ) -> _T_cur: ...

# class RealDictCursor(DictCursorBase):
#     def __init__(self, *args, **kwargs) -> None: ...
#     column_mapping: Any
#     def execute(self, query, vars: Any | None = ...): ...
#     def callproc(self, procname, vars: Any | None = ...): ...
#     def fetchone(self) -> RealDictRow | None: ...  # type: ignore[override]
#     def fetchmany(self, size: int | None = ...) -> list[RealDictRow]: ...  # type: ignore[override]
#     def fetchall(self) -> list[RealDictRow]: ...  # type: ignore[override]
#     def __next__(self) -> RealDictRow: ...  # type: ignore[override]

# class RealDictRow(OrderedDict[Any, Any]):
#     def __init__(self, *args, **kwargs) -> None: ...
#     def __setitem__(self, key, value) -> None: ...

# class NamedTupleConnection(_connection):
#     @overload
#     def cursor(
#         self, name: str | bytes | None = ..., *, withhold: bool = ..., scrollable: bool | None = ...
#     ) -> NamedTupleCursor: ...
#     @overload
#     def cursor(
#         self,
#         name: str | bytes | None = ...,
#         *,
#         cursor_factory: Callable[..., _T_cur],
#         withhold: bool = ...,
#         scrollable: bool | None = ...,
#     ) -> _T_cur: ...
#     @overload
#     def cursor(
#         self, name: str | bytes | None, cursor_factory: Callable[..., _T_cur], withhold: bool = ..., scrollable: bool | None = ...
#     ) -> _T_cur: ...

# class NamedTupleCursor(_cursor):
#     Record: Any
#     MAX_CACHE: int
#     def execute(self, query, vars: Any | None = ...): ...
#     def executemany(self, query, vars): ...
#     def callproc(self, procname, vars: Any | None = ...): ...
#     def fetchone(self) -> NamedTuple | None: ...  # type: ignore[override]
#     def fetchmany(self, size: int | None = ...) -> list[NamedTuple]: ...  # type: ignore[override]
#     def fetchall(self) -> list[NamedTuple]: ...  # type: ignore[override]
#     def __next__(self) -> NamedTuple: ...  # type: ignore[override]

# class LoggingConnection(_connection):
#     log: Any
#     def initialize(self, logobj) -> None: ...
#     def filter(self, msg, curs): ...
#     def cursor(self, *args, **kwargs): ...

# class LoggingCursor(_cursor):
#     def execute(self, query, vars: Any | None = ...): ...
#     def callproc(self, procname, vars: Any | None = ...): ...

# class MinTimeLoggingConnection(LoggingConnection):
#     def initialize(self, logobj, mintime: int = ...) -> None: ...
#     def filter(self, msg, curs): ...
#     def cursor(self, *args, **kwargs): ...

# class MinTimeLoggingCursor(LoggingCursor):
#     timestamp: Any
#     def execute(self, query, vars: Any | None = ...): ...
#     def callproc(self, procname, vars: Any | None = ...): ...

# class LogicalReplicationConnection(_replicationConnection):
#     def __init__(self, *args, **kwargs) -> None: ...

# class PhysicalReplicationConnection(_replicationConnection):
#     def __init__(self, *args, **kwargs) -> None: ...

# class StopReplication(Exception): ...

# class ReplicationCursor(_replicationCursor):
#     def create_replication_slot(self, slot_name, slot_type: Any | None = ..., output_plugin: Any | None = ...) -> None: ...
#     def drop_replication_slot(self, slot_name) -> None: ...
#     def start_replication(
#         self,
#         slot_name: Any | None = ...,
#         slot_type: Any | None = ...,
#         start_lsn: int = ...,
#         timeline: int = ...,
#         options: Any | None = ...,
#         decode: bool = ...,
#         status_interval: int = ...,
#     ) -> None: ...
#     def fileno(self): ...

# class UUID_adapter:
#     def __init__(self, uuid) -> None: ...
#     def __conform__(self, proto): ...
#     def getquoted(self): ...

# def register_uuid(oids: Any | None = ..., conn_or_curs: Any | None = ...): ...

# class Inet:
#     addr: Any
#     def __init__(self, addr) -> None: ...
#     def prepare(self, conn) -> None: ...
#     def getquoted(self): ...
#     def __conform__(self, proto): ...

# def register_inet(oid: Any | None = ..., conn_or_curs: Any | None = ...): ...
# def wait_select(conn) -> None: ...

# class HstoreAdapter:
#     wrapped: Any
#     def __init__(self, wrapped) -> None: ...
#     conn: Any
#     getquoted: Any
#     def prepare(self, conn) -> None: ...
#     @classmethod
#     def parse(cls, s, cur, _bsdec=...): ...
#     @classmethod
#     def parse_unicode(cls, s, cur): ...
#     @classmethod
#     def get_oids(cls, conn_or_curs): ...

# def register_hstore(
#     conn_or_curs, globally: bool = ..., unicode: bool = ..., oid: Any | None = ..., array_oid: Any | None = ...
# ) -> None: ...

# class CompositeCaster:
#     name: Any
#     schema: Any
#     oid: Any
#     array_oid: Any
#     attnames: Any
#     atttypes: Any
#     typecaster: Any
#     array_typecaster: Any
#     def __init__(self, name, oid, attrs, array_oid: Any | None = ..., schema: Any | None = ...) -> None: ...
#     def parse(self, s, curs): ...
#     def make(self, values): ...
#     @classmethod
#     def tokenize(cls, s): ...

# def register_composite(name, conn_or_curs, globally: bool = ..., factory: Any | None = ...): ...
# def execute_batch(cur, sql, argslist, page_size: int = ...) -> None: ...
# def execute_values(cur, sql, argslist, template: Any | None = ..., page_size: int = ..., fetch: bool = ...): ...

# from typing import Any

# from psycopg2._psycopg import (
#     BINARYARRAY as BINARYARRAY,
#     BOOLEAN as BOOLEAN,
#     BOOLEANARRAY as BOOLEANARRAY,
#     BYTES as BYTES,
#     BYTESARRAY as BYTESARRAY,
#     DATE as DATE,
#     DATEARRAY as DATEARRAY,
#     DATETIMEARRAY as DATETIMEARRAY,
#     DECIMAL as DECIMAL,
#     DECIMALARRAY as DECIMALARRAY,
#     FLOAT as FLOAT,
#     FLOATARRAY as FLOATARRAY,
#     INTEGER as INTEGER,
#     INTEGERARRAY as INTEGERARRAY,
#     INTERVAL as INTERVAL,
#     INTERVALARRAY as INTERVALARRAY,
#     LONGINTEGER as LONGINTEGER,
#     LONGINTEGERARRAY as LONGINTEGERARRAY,
#     PYDATE as PYDATE,
#     PYDATEARRAY as PYDATEARRAY,
#     PYDATETIME as PYDATETIME,
#     PYDATETIMEARRAY as PYDATETIMEARRAY,
#     PYDATETIMETZ as PYDATETIMETZ,
#     PYDATETIMETZARRAY as PYDATETIMETZARRAY,
#     PYINTERVAL as PYINTERVAL,
#     PYINTERVALARRAY as PYINTERVALARRAY,
#     PYTIME as PYTIME,
#     PYTIMEARRAY as PYTIMEARRAY,
#     ROWIDARRAY as ROWIDARRAY,
#     STRINGARRAY as STRINGARRAY,
#     TIME as TIME,
#     TIMEARRAY as TIMEARRAY,
#     UNICODE as UNICODE,
#     UNICODEARRAY as UNICODEARRAY,
#     AsIs as AsIs,
#     Binary as Binary,
#     Boolean as Boolean,
#     Column as Column,
#     ConnectionInfo as ConnectionInfo,
#     DateFromPy as DateFromPy,
#     Diagnostics as Diagnostics,
#     Float as Float,
#     Int as Int,
#     IntervalFromPy as IntervalFromPy,
#     ISQLQuote as ISQLQuote,
#     Notify as Notify,
#     QueryCanceledError as QueryCanceledError,
#     QuotedString as QuotedString,
#     TimeFromPy as TimeFromPy,
#     TimestampFromPy as TimestampFromPy,
#     TransactionRollbackError as TransactionRollbackError,
#     Xid as Xid,
#     adapt as adapt,
#     adapters as adapters,
#     binary_types as binary_types,
#     connection as connection,
#     cursor as cursor,
#     encodings as encodings,
#     encrypt_password as encrypt_password,
#     get_wait_callback as get_wait_callback,
#     libpq_version as libpq_version,
#     lobject as lobject,
#     new_array_type as new_array_type,
#     new_type as new_type,
#     parse_dsn as parse_dsn,
#     quote_ident as quote_ident,
#     register_type as register_type,
#     set_wait_callback as set_wait_callback,
#     string_types as string_types,
# )

# ISOLATION_LEVEL_AUTOCOMMIT: int
# ISOLATION_LEVEL_READ_UNCOMMITTED: int
# ISOLATION_LEVEL_READ_COMMITTED: int
# ISOLATION_LEVEL_REPEATABLE_READ: int
# ISOLATION_LEVEL_SERIALIZABLE: int
# ISOLATION_LEVEL_DEFAULT: Any
# STATUS_SETUP: int
# STATUS_READY: int
# STATUS_BEGIN: int
# STATUS_SYNC: int
# STATUS_ASYNC: int
# STATUS_PREPARED: int
# STATUS_IN_TRANSACTION: int
# POLL_OK: int
# POLL_READ: int
# POLL_WRITE: int
# POLL_ERROR: int
# TRANSACTION_STATUS_IDLE: int
# TRANSACTION_STATUS_ACTIVE: int
# TRANSACTION_STATUS_INTRANS: int
# TRANSACTION_STATUS_INERROR: int
# TRANSACTION_STATUS_UNKNOWN: int

# def register_adapter(typ, callable) -> None: ...

# class SQL_IN:
#     def __init__(self, seq) -> None: ...
#     def prepare(self, conn) -> None: ...
#     def getquoted(self): ...

# class NoneAdapter:
#     def __init__(self, obj) -> None: ...
#     def getquoted(self, _null: bytes = ...): ...

# def make_dsn(dsn: Any | None = ..., **kwargs): ...

# JSON: Any
# JSONARRAY: Any
# JSONB: Any
# JSONBARRAY: Any

import psycopg2
import psycopg2.extras
import psycopg2.extensions
import utils.common.log as log



def dict_factory(cursor, row):
        d = {}
        for idx, col in enumerate(cursor.description):
            d[col[0]] = row[idx]
        return d

class PGSqlBase:

    def __init__(self):
        self.conn = None
        self.host = ''
        self.port = 0
        self.user = ''
        self.pwd = ''
        self.dbname = ''

    def __del__(self):
        self.close()

    def connect(self, host, port, user, pwd, dbname):
        """
        连接数据库
        :param host: 
        :param port: 
        :param user: 
        :param pwd: 
        :param dbname: 
        :return: 
        """
        try:
            self.conn = psycopg2.connect(database=dbname, user=user, password=pwd, host=host, port=port)
        except Exception as e:
            log.logger.warning('error: {}'.format(e))
            return False

        self.conn.autocommit = True
        self.host = host
        self.port = port
        self.user = user
        self.pwd = pwd
        self.dbname = dbname
        return True

    def close(self):
        '''
        关闭连接
        :return: 
        '''
        if self.conn is not None:
            self.conn.close()
            self.conn = None

    def _reconnect(self):
        '''重连'''
        self.__del__()
        return self.connect(self.host, self.port, self.user, self.pwd, self.dbname)

    def _exec(self, sql):
        '''执行SQL'''
        ret = 0
        cursor = None
        try:
            cursor = self.conn.cursor()
            # print('[{}]: {}'.format(dt.datetime.now().strftime('%F %T'), sql))
            cursor.execute(sql)
            self.conn.commit()
        #  
        except psycopg2.OperationalError as e:
            log.logger.warning('error: {}'.format(e))
            ret = 2
        except Exception as e:
            log.logger.warning('sql {} , error: {}'.format(sql, e))
            # cursor.execute("rollback")
            ret = 1
        finally:
            if cursor:
                cursor.close()
        return ret

    def exec(self, sql):
       
        b = self._exec(sql)
        if b == 2:
            if self._reconnect():
                b = self._exec(sql)

        if b == 0:
            return True
        else:
            return False


    def _exec_fetchone(self, sql):
  
        dic = {}
        ret = 0
        cursor = None
        try:
            cursor = self.conn.cursor() # 光标
            cursor.execute(sql)
            self.conn.commit()
        #  
        except psycopg2.OperationalError as e:
            log.logger.warning('error: {}'.format(e))
            ret = 2
        except Exception as e:
            log.logger.warning('sql {} , error: {}'.format(sql, e))
            ret = 1
        else:
            r = cursor.fetchone()
            if r is not None:
                dic = {}
                for idx, col in enumerate(cursor.description):
                    dic[col[0]] = r[idx]
        finally:
            if cursor:
                cursor.close()

        return ret, dic


    

    def exec_fetchone(self, sql):
       
        r, dic = self._exec_fetchone(sql)
        if r == 2:
            if self._reconnect():
                r, dic = self._exec_fetchone(sql)

        return dic

    def _exec_fetchall(self, sql):
  
        lnt = []
        ret = 0
        cursor = None
        try:
            cursor = self.conn.cursor()
            cursor.execute(sql)
            self.conn.commit()
        # 
        except psycopg2.OperationalError as e:
            log.logger.warning('error: {}'.format(e))
            ret = 2
        except Exception as e:
            log.logger.warning('sql {} , error: {}'.format(sql, e))
            ret = 1
        else:
            r = cursor.fetchall()
            

            for row in r:
                tmp = {}
                for idx, col in enumerate(cursor.description):
                    tmp[col[0]] = row[idx]
                lnt.append(tmp)
        finally:
            if cursor:
                cursor.close()

        return ret, lnt

    def exec_fetchall(self, sql):
        
        r, lnt = self._exec_fetchall(sql)
        if r == 2:
            if self._reconnect():
                r, lnt = self._exec_fetchall(sql)

        return lnt

    def exec_many(self, sqls):

        b = self._execmany(sqls)
        if b == 2:
            if self._reconnect():
                b = self._execmany(sqls)

        if b == 0:
            return True
        else:
            return False

    def _execmany(self,sqls):

        ret = 0
        cursor = None
        try:
            cursor = self.conn.cursor()
            # print('[{}]: {}'.format(dt.datetime.now().strftime('%F %T'), sql))
            for sql in sqls:
                cursor.execute(sql)
            self.conn.commit()
        #  
        except psycopg2.OperationalError as e:
            log.logger.warning('error: {}'.format(e))
            ret = 2
        except Exception as e:
            log.logger.warning('sql {} , error: {}'.format(sql, e))
            # cursor.execute("rollback")
            ret = 1
        finally:
            if cursor:
                cursor.close()
        return ret